import cv2
import numpy as np
import tkinter as tk
from gui_util import create_root, camera_cap, show_image_on_label
import random
import time
import detect_cube

class ObjClassification:
    def __init__(self):
        self.stable_times = 10
        # 2 char, 1st is trigger (FT), 2nd is color (RGBDU)
        self.states = ['F-']*self.stable_times
        self.arm_action = 1
        self.cap = None
        self.root = None
        self.photo_label = None
        self.video_label = None
        self.last_frame = None
        self.count = [0, 0, 0, 0]  # r, g, b, bad
        self.count_labels = []
        self.class_label = None
        self.button = None
        self.num = 0


    def which_color(part):
        result_rgb = 'D'
        result_rgb=detect_cube.main1()
        if result_rgb == "Rcube":
            print('this is a red cube')
            c='R'
        elif result_rgb == 'Gcube':
            print('this is a gcube cube')
            c='G'
        elif result_rgb == 'Bcube':
            print('this is a blue cube')
            c='B'
        else:
            print('this is a bad cube')
            c='D'
        return c



    def update_count(self, color):
        if color == 'R':
            flag = 0
        elif color == 'G':
            flag = 1
        elif color == 'B':
            flag = 2
        elif color == 'D':  # bad
            flag = 3
            print('bad object ... ')
        else:
            return
        self.count[flag] += 1
        if self.arm_action == 0:
            print('for test only. no arm motion. ')
            return


    def start_click(self):
        self.button = 1


    def exit_click(self):
        self.cap.release()
        time.sleep(1)
        self.root.destroy()

    def take_photo(self):
        show_image_on_label(self.last_frame, self.photo_label, 240, 180)

    def main(self):
        self.root = create_root('./ui_img/class.png')  #设置背景图片class.png
        self.cap = camera_cap()  #获取视频

        photo1 = tk.PhotoImage(file='./ui_img/return.png')  # 设置“返回”按钮的背景图片
        b2 = tk.Button(self.root, image=photo1, width=120, height=50, command=self.exit_click, borderwidth=0, bg='#6ABAFF')
        b2.place(x=1100,y=15)


        photo2 = tk.PhotoImage(file='./ui_img/start_star.png')  # 设置“返回”按钮的背景图片
        b2 = tk.Button(self.root, image=photo2, width=120, height=50, command=self.start_click, borderwidth=0, bg='#6ABAFF')
        b2.place(x=1100,y=85)


        l = tk.Label(self.root, text='0', font=("宋体", 18))  #生成“类别”显示的标签
        self.class_label = l
        l.place(x=990, y=175)  # 将"类别"的标签添加到主窗口

        l = tk.Label(self.root, text='0', font=("Arial", 20)) #生成“红色方块”显示的标签
        self.count_labels.append(l)
        l.place(x=982, y=512)  # 将"红色方块"的标签添加到主窗口的坐标位置

        l = tk.Label(self.root, text='0', font=("Arial", 20)) #生成“绿色方块”显示的标签
        self.count_labels.append(l)
        l.place(x=1128, y=512)  # 将"绿色方块"的标签添加到主窗口

        l = tk.Label(self.root, text='0', font=("Arial", 20)) #生成“蓝色方块”显示的标签
        self.count_labels.append(l)
        l.place(x=982, y=590)  # 将"蓝色方块"的标签添加到主窗口

        l = tk.Label(self.root, text='0', font=("Arial", 20))  #生成“瑕疵品”显示的标签
        self.count_labels.append(l)
        l.place(x=1128, y=590)   # 将"瑕疵品"的标签添加到主窗口  

        l = tk.Label(self.root, text='0', font=("Arial", 20))   #生成“总计”显示的标签
        self.count_labels.append(l)
        l.place(x=1270, y=550)   # 将"总计"的标签添加到主窗口  
   

        self.video_label = tk.Label(self.root)  #生成实时视频框的显示标签
        self.video_label.place(x=184, y=160, height=480, width=640)

        self.photo_label = tk.Label(self.root)  #生成物体结果照片的显示标签
        self.photo_label.place(x=980, y=210, height=180, width=240)

        self.tick_callback()  #执行tick_callback()函数
        self.root.mainloop()

    def tick_callback(self):
        ret, self.last_frame = self.cap.read()
        if not ret:
            print('no frame')
            return
        frame = self.last_frame
        if frame is None: return
        h, w, _ = frame.shape      #去图像的宽和高
        frame = cv2.resize(frame,(w//2,h//2))    #对图像进行缩放
        h,w = h//2, w//2  
        cube = 24
        x1 = (w-cube)//2
        y1 = (h-cube)//2
        cv2.rectangle(frame, (x1,y1), (x1+cube, y1+cube), (0,0,0), 0)    #在图像上画矩形，分别为左上角和右下角坐标，颜色、粗细。

        if self.button == 1:         # 点击按钮开始识别
            time.sleep(0.1)
            # cv2.rectangle(frame, (x1, y1), (x1 + cube, y1 + cube), (0, 0, 0), 0)  # 在图像上画矩形，分别为左上角和右下角坐标，颜色、粗细。
            # part = frame[y1 + 5:y1 + cube - 5, x1 + 5:x1 + cube - 5, :]    #取矩形框内图片
            # part = frame[y1:y1 + cube+25, x1:x1+25 + cube, :]    #取矩形框内图片
            part = frame[0:640, 0:450 + cube, :]    #取矩形框内图片
            cv2.imwrite('inference/images/'+ str(self.num) +'.jpg', part)    #保存到该路径下
            self.num+=1
            c = ObjClassification.which_color(part)    #调用函数进行图像的识别
            # c = 'R'
            t = self.update_state(True, c)
            print("还未进入")
            if t:
                print("以进入")
                part = frame[-25+y1 + 5:35+y1 + cube - 5, -35+x1 + 5:25+x1 + cube - 5, :]  
                part = cv2.rotate(part, cv2.ROTATE_90_COUNTERCLOCKWISE)
                show_image_on_label(part, self.photo_label, 180, 240)
                name = '红色方块' if c == 'R' else (
                    '绿色方块' if c == 'G' else ('蓝色方块' if c == 'B' else ('瑕疵品' if c == 'D' else '未知类别')))
                self.class_label.config(text=name)
                if c != 'U':
                    self.update_count(c)
                    self.show_quantity()
            self.button = 0
        # else:
        #     self.update_state(False, 'U')
        frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
        show_image_on_label(frame, self.video_label, 640, 480)
        self.root.after(30, self.tick_callback)

    def update_state(self, trigger, color='U'):
        # old = self.states[0]
        if not trigger:
            self.states.append('F-')
        else:
            self.states.append('T'+color)
            return True
        # self.states = self.states[-self.stable_times:]
        # if trigger and old != self.states[-1]:
        #     same = all(e == self.states[0] for e in self.states)
        #     if not same:
        #         return True
        return False

    def show_quantity(self):
        print(self.count)
        for i in range(len(self.count)):
            self.count_labels[i].config(text=str(self.count[i]))
        self.count_labels[-1].config(text=str(sum(self.count)))


if __name__ == '__main__':
    ObjClassification().main()

